import numpy as np
import torch
from torch import Tensor

from algorithms.algorithm import Algorithm
from config import device
from individual.population import Population
from operators.crossover.sbx import SBX
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.select_operator import tournament_selection


class NSGA2(Algorithm):
    """
    %------------------------------- Reference --------------------------------
    % K. Deb, A. Pratap, S. Agarwal, and T. Meyarivan, A fast and elitist
    % multiobjective genetic algorithm: NSGA-II, IEEE Transactions on
    % Evolutionary Computation, 2002, 6(2): 182-197.
    """

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs):
        super().__init__(pop_size, problem, kwargs)
        self.population = None
        self.front_no = None
        self.crowd_dis = None
        self.init(pop_size)

    def init(self, pop_size):
        self.population = self.problem.init_population(pop_size)
        _, self.front_no, self.crowd_dis = self.environmental_selection(self.population.obj, self.pop_size, True)

    @staticmethod
    def ND_sort(pop_obj: Tensor, min_sort: int) -> tuple:
        """
        非支配排序
        :param pop_obj: 排序的obj
        :param min_sort: 最少需要排序的个体数，如果只需要得到帕累托排序后的第一层个体，将min_sort设置为1即可
        :return: 返回排序后的结果和最大的帕累托层组成的元组
        """
        pop_obj, loc, repeat_counts = torch.unique(pop_obj, dim=0, return_inverse=True, return_counts=True)
        [row_count, col_count] = pop_obj.shape
        inf = float('inf')
        front_no = torch.full((row_count,), inf, device=device)
        max_front = 0
        while repeat_counts[front_no < inf].sum().item() < min(min_sort, loc.size()[0]):
            max_front += 1
            for i in range(row_count):
                if front_no[i] == inf:
                    dominated = False
                    for j in range(i - 1, -1, -1):
                        if front_no[j] == max_front:
                            m = 1
                            while m < col_count and pop_obj[i, m] >= pop_obj[j, m]:
                                m = m + 1
                            dominated = m == col_count
                            if dominated or col_count == 2:
                                break
                    if not dominated:
                        front_no[i] = max_front;
        front_no = front_no[loc]
        return front_no, max_front

    @staticmethod
    def crowding_distance(obj: Tensor, front_no=None) -> Tensor:
        """
        计算拥挤距离
        :param obj:
        :param front_no:
        :return:
        """
        [obj_size, obj_dim] = obj.shape
        if front_no is None:
            front_no = torch.ones(obj_size, device=device)
        crowd_dis = torch.zeros(obj_size, device=device, dtype=float)
        fronts = torch.tensor(np.setdiff1d(torch.unique(front_no).cpu().numpy(), float('inf')), device=device,
                              dtype=int)
        for i in range(len(fronts)):
            front = torch.nonzero(fronts[i] == front_no, as_tuple=True)[0]
            obj_min = obj[front].min(0)[0].squeeze(dim=0)
            obj_max = obj[front].max(0)[0].squeeze(dim=0)
            for j in range(obj_dim):
                [_, indices] = obj[front][:, j].sort()
                crowd_dis[front[indices[0]]] = float('inf')
                crowd_dis[front[indices[-1]]] = float('inf')
                for k in range(1, len(front) - 1):
                    crowd_dis[front[indices[k]]] = crowd_dis[front[indices[k]]] + (
                            obj[front[indices[k + 1]], j] - obj[front[indices[k - 1]], j]) / (obj_max[j] - obj_min[j])
        return crowd_dis

    @classmethod
    def environmental_selection(cls, obj: Tensor, pop_size: int, return_frontNo_crowdDis: bool = False) -> Tensor:
        """
        环境选择
        :param pop_size:
        :param obj:
        :param return_frontNo_crowdDis: 是否返回front_no，crowd_dis
        :return:next_index,front_no,crowd_dis
        """
        front_no, max_front = cls.ND_sort(obj, pop_size)
        next_index = torch.zeros_like(front_no, device=device, dtype=bool)
        next_index[front_no < max_front] = True
        crowd_dis = cls.crowding_distance(obj, front_no)
        last = torch.nonzero(max_front == front_no, as_tuple=True)[0]
        [_, rank] = crowd_dis[last].sort(descending=True)
        next_index[last[rank[:pop_size - next_index.sum()]]] = True
        return (next_index, front_no, crowd_dis) if return_frontNo_crowdDis else next_index

    def each_iteration(self):
        obj = self.population.obj
        dec = self.population.dec
        mating_pool = tournament_selection(2, self.pop_size, [self.front_no, -self.crowd_dis])
        offspring_dec = SBX.do(dec, dec[mating_pool], self.problem,
                               offspring_size=1)
        offspring_dec = PolynomialMutation.do(offspring_dec, self.problem)
        offspring = self.problem.estimate_population(offspring_dec)
        offspring_obj = offspring.obj
        temp_obj = torch.cat((obj, offspring_obj), dim=0)
        temp_dec = torch.cat((dec, offspring_dec), dim=0)
        next_index, self.front_no, self.crowd_dis = self.environmental_selection(temp_obj, self.pop_size,
                                                                                 return_frontNo_crowdDis=True)
        self.population = Population(temp_dec[next_index], temp_obj[next_index])

    @classmethod
    def get_best(cls, pop_obj, return_front=False):
        front_no, _ = cls.ND_sort(pop_obj, 1)
        return (pop_obj[front_no == 1], front_no) if return_front else pop_obj[front_no == 1]
